package informationretrieval;

import java.util.Collection;
import java.util.HashSet;
import java.util.List;

public class NDCG {

	/**
	 * Private constructor to escape instantiation
	 */
	private NDCG() {
	}

	/**
	 * Calculate the normalized discounted cumulative gain for a result set of
	 * documents
	 * 
	 * @param rankedResultDocs
	 *            The list of the results for a given query
	 * @param truthRank
	 *            The true positive list of documents ID per a given query
	 * @param falseNegative
	 *            The document ID's we should ignore
	 * @return the NDCG for the given list of documents
	 */
	public static double calculateNDCGValue(List<Integer> rankedResultDocs,
			Collection<Integer> truthRank, Collection<Integer> falseNegative) {

		// Method variables
		int rank = 0;
		int falsePositive = 0;
		double result = 0;
		double dcg = 0;
		double idcg = calculateIDCGValue(truthRank.size());

		if (falseNegative == null) {
			falseNegative = new HashSet<Integer>();
		}

		for (int i = 0; i < rankedResultDocs.size(); i++) {
			int docId = rankedResultDocs.get(i);
			if (falseNegative.contains(docId)) {
				falsePositive++;
				continue;
			}

			if (!truthRank.contains(docId)) {
				continue;
			}

			// Get the NDCG value
			rank = i + 1 - falsePositive;
			dcg += (Math.log(2) / Math.log(rank + 1));

		}

		result = (dcg / idcg);

		return result;
	}

	/**
	 * 
	 * Calculates the ideal dcg value for a given result set
	 * 
	 * @param resultSetSize
	 *            The size of the document result set
	 * @return The ideal dcg value for a given set of documents
	 */
	static double calculateIDCGValue(int resultSetSize) {
		double idcg = 0;
		for (int i = 0; i < resultSetSize; i++) {
			idcg += Math.log(2) / Math.log(i + 2);
		}
		return idcg;
	}
}
